"""
本文件主要进行数据集的构建
"""

from torch.utils.data import random_split, Dataset, DataLoader
from config import *
import pandas as pd
import numpy as np
import torch


def train_valid_split(dataset: np.ndarray, valid_ratio: float):
    """ 分割训练集和验证集 """
    valid_size = int(len(dataset) * valid_ratio)
    train_size = len(dataset) - valid_size
    train, valid = random_split(dataset, [train_size, valid_size])
    return np.array(train), np.array(valid)


class CovidDataset(Dataset):
    """ 数据集 """
    def __init__(self, data: np.ndarray, data_flag="train") -> None:
        super().__init__()
        self.data = data
        self.data_flag = data_flag
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return x, y


def train_data_loader():
    """ 返回训练集数据加载器 """
    train_data = pd.read_csv(TRAIN_PATH).values
    train_data, valid_data = train_valid_split(train_data, VALID_RATE)
    train_loader = DataLoader(CovidDataset(train_data, "train"), batch_size=BATCH_SIZE, shuffle=True)
    valid_loader = DataLoader(CovidDataset(valid_data, "train"), batch_size=BATCH_SIZE, shuffle=True)
    return train_loader, valid_loader


def test_data_loader(batch_size=BATCH_SIZE):
    """ 返回测试集数据加载器 """
    test_data = pd.read_csv(TEST_PATH).values
    test_loader = DataLoader(CovidDataset(test_data, "test"), batch_size=batch_size, shuffle=False)
    return test_loader


if __name__ == '__main__':
    # 测试代码
    train_data = pd.read_csv("./data/covid_train.csv").values
    test_data = pd.read_csv("./data/covid_test.csv").values
    # train, valid = train_valid_split(train_data, 0.2)
    train = CovidDataset(train_data, "train")
    test = CovidDataset(test_data, "test")
    print(1)
